import os

import dashscope

# 假设这是你封装好的调用大模型的方法
from dashscope.api_entities.dashscope_response import Message


def save_file(txt, file_name):
    with open(file_name, "w", encoding="utf-8") as f:
        f.write(txt)


def use_llm(system_prompt: str = '你是一位专业的小说作家.', prompt: str = None) -> str:
    """
    模拟调用大语言模型（你可以替换成实际 API 调用）
    """
    print("\n[调用大模型中...]\nPrompt:\n" + prompt)
    messages = [
        Message('system', system_prompt),
        Message('user', prompt),
    ]

    response = dashscope.Generation.call(
        # 若没有配置环境变量，用百炼API Key将下行替换为：api_key="sk-xxx",
        api_key=os.getenv('ai-key'),
        # model="qwen-plus",
        model="qwen-plus-latest",
        # 此处以qwen-plus为例，可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
        messages=messages,
        result_format='message'
    )
    return response.output.choices[0].message['content'].strip()


# 创建项目
def create_project(path, project_name):
    return path


# 根据需求文档生成建表语句
def generate_table(path, requirement_document):
    result = use_llm("根据需求文档生成建表语句", requirement_document)
    file_name = os.path.join(path, 'table.sql')
    save_file(result, file_name)
    return result


# 校验建表语句
def check_table(sql):
    result = use_llm("校验建表语句", sql)
    return result


# 生成entity,mapper,xml,service
def generate_mapper(path, controller, sql):
    result = use_llm("根据建表语句生成entity,mapper,xml,service", sql)
    save_file(result, os.path.join(path, 'mapper.xml'))
    return result


# 生成controller
def generate_controller(path, requirement_document):
    result = use_llm("根据需求文档生成controller", requirement_document)
    return result


# 测试
def test(path, requirement_document):
    result = use_llm("根据需求文档生成测试用例", requirement_document)
    return result
